{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import pandas as pd\r\n",
    "import matplotlib.pyplot as plt\r\n",
    "import numpy as np\r\n",
    "from osgeo import gdal\r\n",
    "gdal.AllRegister()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "%%time\r\n",
    "import sys\r\n",
    "sys.path.append('../03 Slope/')\r\n",
    "import slope as sp\r\n",
    "slopedata = sp.Execution(\"/demcode-learning/dem\", 'Himalaya.tif',extended=True,visible=False)\r\n",
    "slopedata\r\n"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Wall time: 3.92 s\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "array([[28.553305  , 38.62413379, 34.77407285, ..., 12.3688817 ,\n",
       "        14.23310914,  8.47871315],\n",
       "       [43.92916315, 48.24020455, 46.80785695, ..., 12.05719326,\n",
       "        13.94980312,  9.7198677 ],\n",
       "       [45.36775024, 49.07259589, 47.96038106, ..., 10.74441584,\n",
       "        11.99465968,  7.43154972],\n",
       "       ...,\n",
       "       [21.44137841, 21.39522718, 16.05431632, ..., 17.95979969,\n",
       "        29.57329292, 21.56377602],\n",
       "       [20.9236378 , 23.06113562, 20.069195  , ..., 22.22852506,\n",
       "        28.69514066, 19.18076818],\n",
       "       [12.31682409, 18.92602187, 18.80088322, ..., 21.57900759,\n",
       "        24.38870075, 14.02125487]])"
      ]
     },
     "metadata": {},
     "execution_count": 2
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "sys.path.append('../01 ReadDem/')\r\n",
    "import readdem\r\n",
    "info = readdem.read_img('\\demcode-learning\\dem\\Himalaya.tif')\r\n",
    "dem=info[\"data\"]\r\n",
    "dem"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "array([[6167, 6188, 6208, ..., 6074, 6067, 6059],\n",
       "       [6142, 6163, 6182, ..., 6077, 6071, 6063],\n",
       "       [6113, 6135, 6152, ..., 6078, 6073, 6066],\n",
       "       ...,\n",
       "       [5166, 5159, 5153, ..., 5860, 5848, 5830],\n",
       "       [5178, 5170, 5161, ..., 5852, 5838, 5823],\n",
       "       [5188, 5179, 5169, ..., 5842, 5830, 5817]], dtype=int16)"
      ]
     },
     "metadata": {},
     "execution_count": 3
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "ncols = info[\"ncols\"]\r\n",
    "nrows = info[\"nrows\"]"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "# 创建数据类型\r\n",
    "t = np.dtype([('sw', np.float32), \r\n",
    "              ('w', np.float32),\r\n",
    "              ('se', np.float32),\r\n",
    "              ('s', np.float32),\r\n",
    "              ('n', np.float32),\r\n",
    "              ('nw', np.float32),\r\n",
    "              ('ne', np.float32),\r\n",
    "              ('e', np.float32)])\r\n",
    "print(t)\r\n"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[('sw', '<f4'), ('w', '<f4'), ('se', '<f4'), ('s', '<f4'), ('n', '<f4'), ('nw', '<f4'), ('ne', '<f4'), ('e', '<f4')]\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "flow_direction = np.empty([nrows, ncols],dtype=t)\r\n",
    "flow_direction.shape[0]"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "410"
      ]
     },
     "metadata": {},
     "execution_count": 6
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "def Calculatefunctione(gradient):\r\n",
    "    if (gradient < 1):\r\n",
    "        return 8.9*gradient+1.1\r\n",
    "    else:\r\n",
    "        return 10"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "source": [
    "def judgeDirection(directionTuple:list):\r\n",
    "    directionTuple == np.array(directionTuple)\r\n",
    "    if (directionTuple == np.array([1, 0])).all():\r\n",
    "        return {\"dir\":'w',\"dis\":0.5}\r\n",
    "    elif (directionTuple == np.array([0, 1])).all():\r\n",
    "        return {\"dir\": 'n', \"dis\": 0.5}\r\n",
    "    elif (directionTuple == np.array([0, 0])).all():\r\n",
    "        return {\"dir\": 'nw', \"dis\": 0.354}\r\n",
    "    elif (directionTuple == np.array([0, 2])).all():\r\n",
    "        return {\"dir\": 'ne', \"dis\": 0.354}\r\n",
    "    elif (directionTuple == np.array([1, 2])).all():\r\n",
    "        return {\"dir\": 'e', \"dis\": 0.5}\r\n",
    "    elif (directionTuple == np.array([2, 2])).all():\r\n",
    "        return {\"dir\": 'se', \"dis\": 0.354}\r\n",
    "    elif (directionTuple == np.array([2, 1])).all():\r\n",
    "        return {\"dir\": 's', \"dis\": 0.5}\r\n",
    "    elif (directionTuple == np.array([2, 0])).all():\r\n",
    "        return {\"dir\": 'sw', \"dis\": 0.354}\r\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "from numpy.lib import stride_tricks\r\n",
    "# find the cells from neighbors of the dem in 3*3 window by where conditions,index ranges from 1,nrows-2 to 1,ncols-2\r\n",
    "dem_windows = stride_tricks.sliding_window_view(dem, [3, 3])\r\n",
    "slope_windows = stride_tricks.sliding_window_view(slopedata,[3,3])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "source": [
    "value_array = np.empty(shape=[nrows,ncols])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%%timeit\r\n",
    "import math\r\n",
    "for row in range(0,nrows-2):\r\n",
    " for col in range(0,ncols-2):\r\n",
    "    window = dem_windows[row, col]\r\n",
    "    indexs = np.argwhere(window < window[1, 1])\r\n",
    "    if indexs.shape[0] > 0:\r\n",
    "        idx = indexs.shape[0]\r\n",
    "        valueList = []\r\n",
    "        sumvalue = 0\r\n",
    "        for i in range(0, idx):\r\n",
    "            #s_x,s_y value from 0 to 2\r\n",
    "            s_x,s_y = indexs[i]\r\n",
    "            dirc_info = judgeDirection(indexs[i])\r\n",
    "            dirc_str = dirc_info[\"dir\"]\r\n",
    "            dirc_dis = dirc_info[\"dis\"]\r\n",
    "            slope_tan = np.tan(slopedata[row+s_x][col+s_y]*math.pi/180.0)\r\n",
    "            functionE = Calculatefunctione(slope_tan)                \r\n",
    "            value = math.pow(slope_tan, functionE)*dirc_dis\r\n",
    "            sumvalue += value\r\n",
    "            valueList.append({\"dirc_str\": dirc_str, \"value\": value})\r\n",
    "        if sumvalue!=0:\r\n",
    "            for d in valueList:\r\n",
    "                dirc_str = d['dirc_str']\r\n",
    "                flow_direction[row+1][col+1][dirc_str] = d['value']/sumvalue\r\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "np.savetxt(\"result.csv\", flow_direction, fmt='%s')"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python",
   "version": "3.8.10",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.10 64-bit ('geo-env': conda)"
  },
  "interpreter": {
   "hash": "7b0d6b4aa929793d88e84d754c187056723d52cecd68845c7dca810866802715"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}